package cn.edu.sjtu.stap.sootbranchinstr;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import soot.Body;
import soot.BodyTransformer;
import soot.IntType;
import soot.Local;
import soot.PatchingChain;
import soot.RefType;
import soot.Scene;
import soot.SootClass;
import soot.SootMethod;
import soot.Unit;
import soot.Value;
import soot.jimple.AssignStmt;
import soot.jimple.ConditionExpr;
import soot.jimple.IfStmt;
import soot.jimple.InstanceInvokeExpr;
import soot.jimple.IntConstant;
import soot.jimple.InterfaceInvokeExpr;
import soot.jimple.InvokeStmt;
import soot.jimple.Jimple;
import soot.jimple.LookupSwitchStmt;
import soot.jimple.StaticInvokeExpr;
import soot.jimple.Stmt;
import soot.jimple.StringConstant;
import soot.jimple.TableSwitchStmt;
import soot.jimple.VirtualInvokeExpr;
import soot.tagkit.LineNumberTag;
import soot.tagkit.Tag;
import cn.edu.sjtu.stap.sootbranchinstr.template.IfTemplate;
import cn.edu.sjtu.stap.sootbranchinstr.template.PolyTemplate;
import cn.edu.sjtu.stap.sootbranchinstr.template.SwitchTarget;
import cn.edu.sjtu.stap.sootbranchinstr.template.SwitchTemplate;
import cn.edu.sjtu.stap.sootbranchinstr.template.TemplateMaker;

public class BranchInstruTransformer extends BodyTransformer {
	static SootClass recordClass;
	static SootMethod ifRecordMethod;
	static SootMethod switchRecordMethod;
	static SootMethod polyRecordMethod;
	
	static {
		recordClass = Scene.v().loadClassAndSupport("cn.edu.sjtu.stap.sootbranchinstr.BranchRecorder");
		ifRecordMethod = recordClass.getMethod("void recordIfStmt(java.lang.String,java.lang.String,java.lang.String)");
		switchRecordMethod = recordClass.getMethod("void recordSwitchStmt(java.lang.String,java.lang.String,int)");
		polyRecordMethod = recordClass.getMethod("void recordPolyCall(java.lang.String,java.lang.String,java.lang.String,java.lang.String)");
	}
	
	@Override
	protected void internalTransform(Body body, String phase, Map options) {
//		System.out.println(body.getMethod().getSignature());
		PatchingChain<Unit> unitChain = body.getUnits();
		Iterator unitIte = unitChain.snapshotIterator();
		int switchTempCount = 0;
		int virtualCallCount = 0;
		
		Set<Unit> insertedUnitsInMethod = new HashSet<Unit>();
		
		while(unitIte.hasNext()){
			Unit unit = (Unit)unitIte.next();
			if(unit.getTag("LineNumberTag") == null) continue;
			
				// Here we consider three kinds of jumps in Jimple:
				// if, lookupswitch, and tableswitch
				if(unit instanceof IfStmt && !insertedUnitsInMethod.contains(unit) ){
					String origLineNumber = unit.getTag("LineNumberTag").toString();
					IfStmt ifStmt = (IfStmt)unit;
					Value cond = ifStmt.getCondition();
					
					if(cond instanceof ConditionExpr){
						// the argument list of the method call to record branch taken event
						List argList1 = new ArrayList();
						argList1.add(StringConstant.v(body.getMethod().getSignature()));
						argList1.add(StringConstant.v(ifStmt.getTag("LineNumberTag").toString()));
						argList1.add(StringConstant.v("taken"));
						
						StaticInvokeExpr invokeExp = 
							Jimple.v().newStaticInvokeExpr(ifRecordMethod.makeRef(), argList1);
						Stmt recordTaken = Jimple.v().newInvokeStmt(invokeExp);
						unitChain.insertBefore(recordTaken, unit);
//						insertedUnitsInMethod.add(recordTaken);
						recordTaken.addTag(new LineNumberTag(Integer.valueOf(origLineNumber)));
						
						Stmt gotoOrigIfStmt = Jimple.v().newGotoStmt(ifStmt);
						unitChain.insertBefore(gotoOrigIfStmt, recordTaken);
//						insertedUnitsInMethod.add(gotoOrigIfStmt);
						gotoOrigIfStmt.addTag(new LineNumberTag(Integer.valueOf(origLineNumber)));
						
						// the argument list of the method call to record branch untaken event
						List argList2 = new ArrayList();
						argList2.add(StringConstant.v(body.getMethod().getSignature()));
						argList2.add(StringConstant.v(ifStmt.getTag("LineNumberTag").toString()));
						argList2.add(StringConstant.v("untaken"));
						StaticInvokeExpr invokeExp2 = 
							Jimple.v().newStaticInvokeExpr(ifRecordMethod.makeRef(), argList2);
						Stmt recordFallThrough = Jimple.v().newInvokeStmt(invokeExp2);
						unitChain.insertBefore(recordFallThrough, gotoOrigIfStmt);
//						insertedUnitsInMethod.add(recordFallThrough);
						recordFallThrough.addTag(new LineNumberTag(Integer.valueOf(origLineNumber)));
						
						IfStmt synIfStmt = Jimple.v().newIfStmt(cond, recordTaken);
						unitChain.insertBefore(synIfStmt, recordFallThrough);
						insertedUnitsInMethod.add(synIfStmt);
						synIfStmt.addTag(new LineNumberTag(Integer.valueOf(origLineNumber)));
						
						// make a template for this if statement
						IfTemplate it = new IfTemplate();
						it.setKey(body.getMethod().getSignature() + BranchRecorder.DATA_SEPARATOR + ifStmt.getTag("LineNumberTag").toString());
						it.setTakenTarget(ifStmt.getTarget().getTag("LineNumberTag").toString());
						TemplateMaker.addIfTemplate(it);
					}
				}else if(unit instanceof TableSwitchStmt){
					String origLineNumber = unit.getTag("LineNumberTag").toString();
					TableSwitchStmt tsStmt = (TableSwitchStmt)unit;
					
					List argList3 = new ArrayList();
					argList3.add(StringConstant.v(body.getMethod().getSignature()));
					argList3.add(StringConstant.v(unit.getTag("LineNumberTag").toString()));
					
					// Here is some risk of local variable redefinition
					Local switchLocal = Jimple.v().newLocal("st_" + (switchTempCount++), IntType.v());
					body.getLocals().add(switchLocal);
					AssignStmt localAssign = Jimple.v().newAssignStmt(switchLocal, tsStmt.getKey());
					unitChain.insertBefore(localAssign, tsStmt);
					localAssign.addTag(new LineNumberTag(Integer.valueOf(origLineNumber)));
					
					argList3.add(switchLocal);
					StaticInvokeExpr invokeExp3 = Jimple.v().newStaticInvokeExpr(switchRecordMethod.makeRef(), argList3);
					Stmt recordSwitch = Jimple.v().newInvokeStmt(invokeExp3);
					unitChain.insertAfter(recordSwitch, localAssign);
					recordSwitch.addTag(new LineNumberTag(Integer.valueOf(origLineNumber)));
					
					// make a template for this switch statement
					int lowIndex = tsStmt.getLowIndex();
					int highIndex = tsStmt.getHighIndex();
					SwitchTemplate st = new SwitchTemplate();
					st.setKey(body.getMethod().getSignature() + BranchRecorder.DATA_SEPARATOR + unit.getTag("LineNumberTag").toString());
					List<SwitchTarget> targetList = new ArrayList<SwitchTarget>();
					for(int i = lowIndex, j = 0; i <= highIndex; i++, j++){
						SwitchTarget t = new SwitchTarget();
						t.setTargetIndex(String.valueOf(i));
						t.setTargetLine(tsStmt.getTarget(j).getTag("LineNumberTag").toString());
						targetList.add(t);
					}
					
					SwitchTarget defaultTarget = new SwitchTarget();
					defaultTarget.setTargetIndex("default");
					defaultTarget.setTargetLine(tsStmt.getDefaultTarget().getTag("LineNumberTag").toString());
					targetList.add(defaultTarget);
					st.setTargets(targetList);
					TemplateMaker.addSwitchTemplate(st);
				}else if(unit instanceof LookupSwitchStmt){
					String origLineNumber = unit.getTag("LineNumberTag").toString();
					LookupSwitchStmt lsStmt = (LookupSwitchStmt)unit;
					
					List argList4 = new ArrayList();
					argList4.add(StringConstant.v(body.getMethod().getSignature()));
					argList4.add(StringConstant.v(unit.getTag("LineNumberTag").toString()));
					
					// Here is some risk of local variable redefinition
					Local switchLocal = Jimple.v().newLocal("st_" + (switchTempCount++), IntType.v());
					body.getLocals().add(switchLocal);
					AssignStmt localAssign = Jimple.v().newAssignStmt(switchLocal, lsStmt.getKey());
					unitChain.insertBefore(localAssign, lsStmt);
					localAssign.addTag(new LineNumberTag(Integer.valueOf(origLineNumber)));
					
					argList4.add(switchLocal);
					
					StaticInvokeExpr invokeExp4 = Jimple.v().newStaticInvokeExpr(switchRecordMethod.makeRef(), argList4);
					Stmt recordSwitch = Jimple.v().newInvokeStmt(invokeExp4);
					unitChain.insertAfter(recordSwitch, localAssign);
					recordSwitch.addTag(new LineNumberTag(Integer.valueOf(origLineNumber)));
					
					// make a template for this switch statement
					SwitchTemplate st = new SwitchTemplate();
					st.setKey(body.getMethod().getSignature() + BranchRecorder.DATA_SEPARATOR + unit.getTag("LineNumberTag").toString());
					List<SwitchTarget> tList = new ArrayList<SwitchTarget>();
					List indexList = lsStmt.getLookupValues();
					int targetCounter = 0;
					for(Object value : indexList){
						IntConstant icValue = (IntConstant)value;
						Unit tu = lsStmt.getTarget(targetCounter++);
						SwitchTarget t = new SwitchTarget();
						t.setTargetIndex(String.valueOf(icValue.value));
						t.setTargetLine(tu.getTag("LineNumberTag").toString());
						tList.add(t);
					}
					
					SwitchTarget t = new SwitchTarget();
					t.setTargetIndex("default");
					t.setTargetLine(lsStmt.getDefaultTarget().getTag("LineNumberTag").toString());
					tList.add(t);
					st.setTargets(tList);
					TemplateMaker.addSwitchTemplate(st);
				}else if(unit instanceof InvokeStmt && !insertedUnitsInMethod.contains(unit)){
					String origLineNumber = unit.getTag("LineNumberTag").toString();
					InvokeStmt is = (InvokeStmt)unit;
					
					// Here we do not consider call to constructors (i.e., specialinvoke)
					if(   (is.getInvokeExpr() instanceof InstanceInvokeExpr)
						&&	((is.getInvokeExpr() instanceof InterfaceInvokeExpr)
					  || (is.getInvokeExpr() instanceof VirtualInvokeExpr))  // here virtualinvoke has different meanings from the virtual call.
					){
						// TODO: enable this condition check after testing phase.
						if(VirtualCallDatabase.isVirtualCallExpression(body.getMethod().getSignature(), is)){
							int currentCount = virtualCallCount++;
							
							Local classRef = Jimple.v().newLocal("cn_" + currentCount, RefType.v("java.lang.Class"));
							body.getLocals().add(classRef);
							Local baseRef = Jimple.v().newLocal("vb_" + currentCount, RefType.v("java.lang.Object"));
							body.getLocals().add(baseRef);
							AssignStmt baseAssign = Jimple.v().newAssignStmt(baseRef, ((InstanceInvokeExpr)is.getInvokeExpr()).getBase());
							unitChain.insertBefore(baseAssign, unit);
							baseAssign.addTag(new LineNumberTag(Integer.valueOf(origLineNumber)));
//							insertedUnitsInMethod.add(baseAssign);
							
							SootMethod getClassMethod = Scene.v().getMethod("<java.lang.Object: java.lang.Class getClass()>");
							VirtualInvokeExpr getClassInvoke = Jimple.v().newVirtualInvokeExpr(baseRef, getClassMethod.makeRef());
							AssignStmt classAssign = Jimple.v().newAssignStmt(classRef, getClassInvoke);
							unitChain.insertAfter(classAssign, baseAssign);
							classAssign.addTag(new LineNumberTag(Integer.valueOf(origLineNumber)));
//							insertedUnitsInMethod.add(classAssign);
							
							SootMethod getNameMethod = Scene.v().getMethod("<java.lang.Class: java.lang.String getName()>");
							VirtualInvokeExpr getNameInvoke = Jimple.v().newVirtualInvokeExpr(classRef, getNameMethod.makeRef());
							Local classNameLocal = Jimple.v().newLocal("vc_" + currentCount, RefType.v("java.lang.String"));
							body.getLocals().add(classNameLocal);
							AssignStmt nameAssign = Jimple.v().newAssignStmt(classNameLocal, getNameInvoke);
							unitChain.insertAfter(nameAssign, classAssign);
							nameAssign.addTag(new LineNumberTag(Integer.valueOf(origLineNumber)));
//							insertedUnitsInMethod.add(nameAssign);
							
							List argList5 = new ArrayList();
							argList5.add(StringConstant.v(body.getMethod().getSignature()));
							argList5.add(StringConstant.v(unit.getTag("LineNumberTag").toString()));
							argList5.add(classNameLocal);
							argList5.add(StringConstant.v(is.getInvokeExpr().getMethod().getSubSignature()));
							
							StaticInvokeExpr invokeExp5 = Jimple.v().newStaticInvokeExpr(polyRecordMethod.makeRef(), argList5);
							Stmt recordPoly = Jimple.v().newInvokeStmt(invokeExp5);
							unitChain.insertAfter(recordPoly, nameAssign);
							recordPoly.addTag(new LineNumberTag(Integer.valueOf(origLineNumber)));
//							insertedUnitsInMethod.add(recordPoly);
							
							// make a template for this polymorphic method call
							PolyTemplate pt = new PolyTemplate();
							pt.setKey(body.getMethod().getSignature() 
									+ BranchRecorder.DATA_SEPARATOR + unit.getTag("LineNumberTag").toString()
									+ BranchRecorder.DATA_SEPARATOR + is.getInvokeExpr().getMethod().getSubSignature()
							);
							pt.setCallType(((InstanceInvokeExpr)is.getInvokeExpr()).getBase().getType().toString());
							
							List<String> tList = new ArrayList<String>();
//							tList.add(is.getInvokeExpr().getMethod().getSignature());
							tList.add(VirtualCallDatabase.calleeSigOfVirtualEdge);
							for(String optionMethodSig : VirtualCallDatabase.optionMethodSigs){
								tList.add(optionMethodSig);
							}
							pt.setTargets(tList);
							TemplateMaker.addPolyTemplate(pt);
						}
					}
				}
			
		}
	}
}
